# -*- coding: utf-8 -*-
import json
import base64
from PIL import Image
import io
import sys
import websocket
import threading
import numpy as np
import easyocr

ws = websocket.create_connection('ws://localhost:8081')
print(ws)
ocr_engine = easyocr.Reader(['ch_sim', 'en'])
ocr_lock = threading.Lock()  # 用于同步访问 OCR 引擎

def imgOcr(path_or_base64):
    global ocr_engine
    with ocr_lock:
        # 解码和图像处理部分无需加锁
        image_data = base64.b64decode(
            path_or_base64.split(',')[1] if path_or_base64.startswith('data:image') else path_or_base64)


        with io.BytesIO(image_data) as buffer, Image.open(buffer) as image:
            # 进行裁剪，获取需要的区域
            temp_image = image.crop((0, 712, 1080, 1500))

            userName_image_np = np.array(temp_image)
            ocr_results = ocr_engine.readtext(userName_image_np, detail=0)
            print(ocr_results)

def userNameOcr(path_or_base64):
    global ocr_engine
    with ocr_lock:
        if path_or_base64.startswith('data:image'):
            # 去除 base64 字符串中的头部信息（例如 'data:image/jpeg;base64,'）
            path_or_base64 = path_or_base64.split(',')[1]
        # 打开图片
        # image = Image.open(path)
        image_data = base64.b64decode(path_or_base64)  # 解码 base64 字符串
        image = Image.open(io.BytesIO(image_data))  # 将字节数据转换为图片
        # 截取用户姓名和用户信息区域
        userName_image = image.crop((0, 2070, 1080, 2200))

        # 将 PIL 图像转换为 numpy 数组
        userName_image_np = np.array(userName_image)
        ocr_results = ocr_engine.readtext(userName_image_np, detail=0)
        print(ocr_results)
        user = ''.join(ocr_results[0:len(ocr_results) - 1])

        if len(user)<= 0:
            return True
        else:
            return False

def handle_message(image_message):
    imgOcr(image_message)
    # if ws.connected:
    #     ws.send(json.dumps(image_message))

def main():
    while True:
        try:
            result = ws.recv()
        except Exception as e:
            ws.close()
            print("WebSocket 连接已关闭。")
            sys.exit()

        handle_message(result)

if __name__ == '__main__':
    main()
